#include <iostream>
#include <vector>
#include <map>

#include <poll.h>
#include <sys/stat.h>
#include <sys/select.h>
#include <sys/epoll.h>

#include "../common/filehelper.hpp"

class IOHelper {
protected:
    bool add_file_desc(int fd, FileDesc *desc, bool force=false) {
        auto it = m_file_map.find(fd);
        if (it != m_file_map.end() && !force) {
            return false;
        }

        m_file_map[fd] = desc;

        return true;
    }

    bool remove_file_desc_from_fd(int fd) {
        try {
            m_file_map.erase(fd);
            return true;
        } catch (const std::exception &e) {
            return false;
        }
    }

public:
    FileDesc *get_file_from_fd(int fd) {
        try {
            return m_file_map[fd];
        } catch (const std::exception &e) {
            return NULL;
        }
    }

protected:
    std::map<int, FileDesc *> m_file_map;

};

class PollHelper: public IOHelper {
public:
    PollHelper(void) = default;
    ~PollHelper(void) = default;

    bool add_file(FileDesc *f, short int events);

    bool remove_file(int fd);

    int poll(int timeout);

    std::vector<struct pollfd> get_pollfds(void) {
        return m_pfds;
    }

private:
    std::vector<struct pollfd> m_pfds;

};

class SelectHelper: public IOHelper {
public:

    bool add_file(FileDesc *f, bool read, bool write, bool except=false);
    bool remove_file(int fd);

    int select(int timeout) {
        struct timeval tv;

        if (timeout > 0) {
            tv.tv_sec = timeout / 1000;
            tv.tv_usec = (timeout % 1000) * 1000;
        }

        return SelectHelper::select(timeout < 0 ? NULL : &tv);
    }

    int select(struct timeval *tv);

    const std::vector<int> &get_fds(void) const {
        return m_fds;
    }

    const fd_set &get_readfds(void) const {
        return m_readfds;
    }

    const fd_set &get_writefds(void) const {
        return m_writefds;
    }

    const fd_set &get_exceptfds(void) const {
        return m_exceptfds;
    }

private:
    std::vector<int> m_fds;
    fd_set m_readfds;
    fd_set m_writefds;
    fd_set m_exceptfds;
    int m_fd_max {0};
};

#include <sys/epoll.h>
class EpollHelper: public IOHelper {
public:
    EpollHelper(int flags=0) {
        m_epollfd = epoll_create1(flags);
    }

    ~EpollHelper() {
        if (m_epollfd >= 0) {
            close(m_epollfd);
        }
    }

    bool is_ok(void) const {
        return m_epollfd >= 0;
    }

    bool add_epoll(FileDesc *f, uint32_t events) {
        struct epoll_event ev = {
            .events = events,
            .data = {
                .fd = f->get_fd(),
            },
        };

        return EpollHelper::add_epoll(f, &ev);
    }

    bool modify_epoll(FileDesc *f, uint32_t events) {
        struct epoll_event ev = {
            .events = events,
        };

        return EpollHelper::modify_epoll(f, &ev);
    }

    bool add_epoll(FileDesc *f, const struct epoll_event *event);
    bool modify_epoll(FileDesc *f, const struct epoll_event *event);
    // Since Linux 2.6.9, event can be specified as NULL when using EPOLL_CTL_DEL.
    bool remove_epoll(FileDesc *f, const struct epoll_event *event=NULL);


    int wait(std::vector<struct epoll_event> &events, int timeout);   // milliseconds

private:
    std::vector<int> m_fds;
    int m_epollfd {-1};
};
